import base64

import fastdeploy as fd
import cv2
import os

from api.src.db import repositories
from config import appSettings


def build_option(device: str, use_trt: bool):
    option = fd.RuntimeOption()
    if device.lower() == "gpu":
        option.use_gpu()
        option.use_paddle_infer_backend()

    if use_trt:
        option.use_trt_backend()
        # If use original Tensorrt, not Paddle-TensorRT,
        # comment the following two lines
        option.enable_paddle_to_trt()
        option.enable_paddle_trt_collect_shape()
        option.set_trt_input_shape("img", [1, 3, 512, 512])

    if device.lower() == "kunlunxin":
        option.use_kunlunxin()
    return option


def generate_filename(path: str, ext: str):
    dirname, filename = os.path.split(path)
    filename_arr = filename.split(".")
    file_name = filename_arr[0]
    return file_name + ext


def base64_encode_img(img):
    with open(img, "rb") as f:
        return base64.b64encode(f.read())


class AIService:
    def __init__(
        self, model_path: str = None, image_input: str = None, save_path: str = None
    ):
        self.model_path = model_path
        self.image_input = image_input
        self.save_path = save_path
        self.runtime_option = build_option("cpu", False)

        return

    def setup_runtime(self):
        model_file = os.path.join(self.model_path, "model.pdmodel")
        params_file = os.path.join(self.model_path, "model.pdiparams")
        config_file = os.path.join(self.model_path, "deploy.yaml")
        model = fd.vision.matting.PPMatting(
            model_file, params_file, config_file, runtime_option=self.runtime_option
        )
        return model

    def predict(self, image_path: str):
        im = cv2.imread(image_path)
        model = self.setup_runtime()
        return model.predict(im), im

    def remove_bg(self, image_path: str):
        fg_name = generate_filename(image_path, "_fg.png")
        fg_file = self.save_path + "/" + fg_name
        res, im = self.predict(image_path)
        vis_im = fd.vision.vis_matting(im, res, True, 0.5, True)
        # bg_transparent = cv2.imread(appSettings.picture_bg_path + "/white_t.png")
        # vis_im_with_bg = fd.vision.swap_background(im, bg_transparent, res)
        cv2.imwrite(fg_file, vis_im)
        resReturn = {"alpha": base64_encode_img(fg_file)}
        return resReturn

    def replace_bg(self, image_path: str, bg_image_color: str):
        fg_name = generate_filename(image_path, "_fg.jpg")
        fg_file = self.save_path + "/" + fg_name
        res, im = self.predict(image_path)
        vis_im = fd.vision.vis_matting(im, res, True)
        cv2.imwrite(fg_file, vis_im)
        resReturn = {"red": "", "white": "", "blue": ""}
        if bg_image_color == "all":
            bg_name_red = generate_filename(image_path, f"_red.jpg")
            bg_name_blue = generate_filename(image_path, f"_blue.jpg")
            bg_name_white = generate_filename(image_path, f"_white.jpg")

            bg_image_red_path = os.path.join(self.image_input, "red.jpg")
            bg_red = cv2.imread(bg_image_red_path)
            bg_image_blue_path = os.path.join(self.image_input, "blue.jpg")
            bg_blue = cv2.imread(bg_image_blue_path)
            bg_image_white_path = os.path.join(self.image_input, "white.jpg")
            bg_white = cv2.imread(bg_image_white_path)

            vis_im_with_red_bg = fd.vision.swap_background(im, bg_red, res)
            vis_im_with_blue_bg = fd.vision.swap_background(im, bg_blue, res)
            vis_im_with_white_bg = fd.vision.swap_background(im, bg_white, res)
            cv2.imwrite(self.save_path + "/" + bg_name_red, vis_im_with_red_bg)
            cv2.imwrite(self.save_path + "/" + bg_name_blue, vis_im_with_blue_bg)
            cv2.imwrite(self.save_path + "/" + bg_name_white, vis_im_with_white_bg)
            resReturn["red"] = base64_encode_img(self.save_path + "/" + bg_name_red)
            resReturn["blue"] = base64_encode_img(self.save_path + "/" + bg_name_blue)
            resReturn["white"] = base64_encode_img(self.save_path + "/" + bg_name_white)
        else:
            bg_name = generate_filename(image_path, f"_{bg_image_color}.jpg")
            bg_image_path = os.path.join(self.image_input, f"{bg_image_color}.jpg")
            bg = cv2.imread(bg_image_path)
            vis_im_with_bg = fd.vision.swap_background(im, bg, res)
            cv2.imwrite(self.save_path + "/" + bg_name, vis_im_with_bg)
            resReturn[bg_image_color] = base64_encode_img(
                self.save_path + "/" + bg_name
            )
        return resReturn
